#!/bin/bash

# 检查系统是否为 Ubuntu 18.04 或其他支持的版本
distribution=$(. /etc/os-release; echo $ID$VERSION_ID)
if [ -z "$distribution" ]; then
  echo "无法识别操作系统版本，确保您的系统为 Ubuntu 18.04 或更高版本。"
  exit 1
fi

# 提示信息并确认
echo "即将停止所有正在运行的 Docker 容器，安装 nvidia-driver-470 和 nvidia-docker 并重启 Docker 服务。"
read -p "确定继续吗？ (y/n): " confirm
if [[ "$confirm" != "y" ]]; then
  echo "操作已取消。"
  exit 1
fi

# 记录所有当前运行的容器ID
echo "获取当前运行的容器ID..."
running_containers=$(docker ps -q)

# 停止所有正在运行的容器
if [ -n "$running_containers" ]; then
  echo "正在停止所有容器..."
  docker stop $running_containers
  echo "所有容器已停止。"
else
  echo "没有正在运行的容器。"
fi

# 添加 NVIDIA PPA 仓库
echo "添加 NVIDIA PPA 仓库..."
sudo add-apt-repository -y ppa:graphics-drivers/ppa
sudo apt-get update

# 安装 nvidia-driver-470
echo "安装 nvidia-driver-470..."
sudo apt-get install -y nvidia-driver-470

# 自动删除旧版本的 NVIDIA 驱动程序
echo "清理旧版本的 NVIDIA 驱动程序..."
sudo apt-get autoremove -y

# 添加 nvidia-docker 仓库的 GPG key 和仓库地址
echo "添加 nvidia-docker GPG key..."
sudo apt-key adv --fetch-keys https://nvidia.github.io/nvidia-docker/gpgkey
echo "添加 nvidia-docker 仓库地址..."
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

# 更新 apt 软件包列表
echo "更新软件包列表..."
sudo apt-get update

# 安装 nvidia-docker 和相关依赖
echo "安装 nvidia-container-toolkit、nvidia-docker2 和 nvidia-cuda-toolkit..."
sudo apt-get install -y nvidia-container-toolkit nvidia-docker2 nvidia-cuda-toolkit

# 备份并替换 daemon.json 文件
if sudo test -f /etc/docker/daemon.json; then
  echo "备份当前 daemon.json 文件..."
  sudo mv /etc/docker/daemon.json /etc/docker/daemon.json.back
fi
echo "替换 daemon.json 文件..."
sudo cp ./daemon.json /etc/docker/daemon.json

# 重启 Docker 服务
echo "重启 Docker 服务..."
sudo systemctl restart docker

# 检查 Docker 服务状态
if systemctl is-active --quiet docker; then
  echo "Docker 服务已成功重启。"
else
  echo "Docker 服务重启失败，请检查系统日志。"
  exit 1
fi

# 重新启动之前关闭的容器
if [ -n "$running_containers" ]; then
  echo "重新启动之前停止的容器..."
  docker start $running_containers
  echo "所有容器已重新启动。"
else
  echo "没有容器需要重新启动。"
fi

# 提示用户是否确认继续重启宿主机
echo "nvidia-driver-470 与 nvidia-docker 的安装配置已完成！"
read -p "是否立即重启宿主机以应用更改？ (y/n): " reboot_confirm
if [[ "$reboot_confirm" == "y" ]]; then
  # 立即重启宿主机
  echo "正在重启宿主机..."
  sudo reboot
else
  echo "操作已取消。请手动重启宿主机以应用更改。"
fi
